package com.ping;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.SocketTimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

/**
 * This class sends a ping to the remote side when timeout on reading. If no
 * failure is detected, it retries until at least a byte is read.
 */
public class PingInputStream extends FilterInputStream {
	private DataOutputStream out;
	private int rpcTimeout;
	private ByteArrayOutputStream pingRequest; // ping message
	private AtomicBoolean running = new AtomicBoolean(true); // if client runs
	private AtomicLong lastActivity = new AtomicLong();// last I/O activity time
	private AtomicBoolean shouldCloseConnection = new AtomicBoolean(); // indicate if the connection is closed
	private int pingInterval; // how often sends ping to the server in msecs
	/* constructor */
	protected PingInputStream(InputStream in) {
		super(in);
	}

	/*
	 * Process timeout exception if the connection is not going to be closed or
	 * is not configured to have a RPC timeout, send a ping. (if rpcTimeout is
	 * not set to be 0, then RPC should timeout. otherwise, throw the timeout
	 * exception.
	 */
	private void handleTimeout(SocketTimeoutException e) throws IOException {
		if (shouldCloseConnection.get() || !running.get() || rpcTimeout > 0) {
			throw e;
		} else {
			sendPing();
		}
	}

	/**
	 * Read a byte from the stream. Send a ping if timeout on read. Retries if
	 * no failure is detected until a byte is read.
	 * 
	 * @throws IOException
	 *             for any IO problem other than socket timeout
	 */
	@Override
	public int read() throws IOException {
		do {
			try {
				return super.read();
			} catch (SocketTimeoutException e) {
				handleTimeout(e);
			}
		} while (true);
	}

	/**
	 * Read bytes into a buffer starting from offset <code>off</code> Send a
	 * ping if timeout on read. Retries if no failure is detected until a byte
	 * is read.
	 * 
	 * @return the total number of bytes read; -1 if the connection is closed.
	 */
	@Override
	public int read(byte[] buf, int off, int len) throws IOException {
		do {
			try {
				return super.read(buf, off, len);
			} catch (SocketTimeoutException e) {
				handleTimeout(e);
			}
		} while (true);
	}

	/*
	 * Send a ping to the server if the time elapsed since last I/O activity is
	 * equal to or greater than the ping interval
	 */
	private synchronized void sendPing() throws IOException {
		long curTime = Time.now();
		if (curTime - lastActivity.get() >= pingInterval) {
			lastActivity.set(curTime);
			synchronized (out) {
				out.writeInt(pingRequest.size());
				pingRequest.writeTo(out);
				out.flush();
			}
		}
	}
}